from pydantic import BaseModel, Field
from langchain_core.prompts import PromptTemplate
from langchain_ollama import OllamaLLM
from langchain_core.output_parsers import PydanticOutputParser
import re


# 定义输出解析器的自定义数据结构
class Answer(BaseModel):
    question: str = Field(description="用户提出的问题")
    answer: str = Field(description="对用户问题做出的简洁回答")
    details: str = Field(description="对用户问题做出的详细回答")


class PromptEngineerModule:
    def __init__(self):
        # 使用本地Ollama的Qwen3:4B模型
        self.llm = OllamaLLM(
            model="qwen3:4b",
            temperature=0.3,
            num_ctx=4096
        )
        # 创建输出解析器
        self.parser = PydanticOutputParser(pydantic_object=Answer)

        # 创建提示模板
        # self.prompt_template = PromptTemplate.from_template(
        #     "请用中文回答以下问题，从知识库中查找，并提供简洁答案和详细解释，如果没找到就直接告诉我没有：\n\n"
        #     "问题：{question}\n\n"
        #     "请严格按以下JSON格式输出，不要包含任何其他信息：\n"
        #     "{format_instructions}\n\n"
        #     "注意：\n"
        #     "1. 简洁回答(answer)控制在50字以内\n"
        #     "2. 详细解释(details)需要包含技术细节\n"
        #     "3. 确保JSON格式正确"
        # ).partial(
        #     format_instructions=self.parser.get_format_instructions()
        # )

    def create_prompt_template (self,query:str)->str:
        prompt_template = PromptTemplate(
            template=(
        "请用中文回答以下问题，从知识库中查找，并提供简洁答案和详细解释，如果没找到就直接告诉我没有：\n\n"
        "问题：{question}\n\n"
        "请严格按以下JSON格式输出，不要包含任何其他信息：\n"
        "{format_instructions}\n\n"
        "注意：\n"
        "1. 简洁回答(answer)控制在50字以内\n"
        "2. 详细解释(details)需要包含技术细节\n"
        "3. 确保JSON格式正确"
        ),
        input_variables=["question"],
        partial_variables={
            "format_instructions": self.parser.get_format_instructions()}
        )

        return prompt_template.format(question=query)

    # def create_prompt(self, query: str) -> str:
    #     """创建格式化的提示词"""
    #     return self.prompt_template.format(question=query)

    def parse_response(self, response: str) -> Answer:
        """解析模型响应为结构化数据"""
        try:
            # 尝试直接解析
            return self.parser.parse(response)
        except Exception as e:
            print(f"解析失败: {str(e)}")
            print(f"原始响应: {response}")

            # 尝试清理响应
            cleaned = re.sub(r'```json|```', '', response).strip()
            try:
                return self.parser.parse(cleaned)
            except:
                # 尝试提取JSON部分
                json_match = re.search(r'\{[\s\S]*\}', response)
                if json_match:
                    return self.parser.parse(json_match.group(0))
                raise ValueError("无法解析模型响应")

    def generate_answer(self, query: str) -> Answer:
        """生成结构化答案"""
        # 创建提示词
        prompt = self.create_prompt_template(query)
        print(f"生成的提示词:\n\n{prompt}\n")

        # 调用本地模型 (使用invoke方法)
        response = self.llm(prompt)
        print(f"\n原始模型响应:\n\n{response}\n")

        # 解析响应
        parsed_answer = self.parse_response(response)

        # 验证解析结果
        if not parsed_answer.question or not parsed_answer.answer:
            raise ValueError("解析结果缺少必要字段")

        # 打印结构化结果
        print("\n解析后的结构化结果:")
        print(f"问题: {parsed_answer.question}")
        print(f"简洁回答: {parsed_answer.answer}")
        print(f"详细解释: {parsed_answer.details}")

        return parsed_answer


if __name__ == "__main__":
    print("正在初始化本地Qwen3:4B模型...")
    prompt_engineering = PromptEngineerModule()

    query = "无人机端需要小模型做实时处理用什么板子"
    parse_response = prompt_engineering.generate_answer(query)
